﻿namespace Cellfish.Toolbox.Tasks
{
    using System;
    using System.Threading;
    using System.Threading.Tasks;

    /// <summary>
    /// Extension methods to add a timeout to executing tasks.
    /// </summary>
    public static class WithTimeoutExtensions
    {
        /// <summary>
        /// Adds a timeout to your task. If task does not complete before timeout then 
        /// the task fails with a TimeoutException.
        /// </summary>
        /// <remarks>Note that the task is not automatically cancelled if the timeout occurs.</remarks>
        /// <param name="task">The task that should have a timeout added.</param>
        /// <param name="timeout">The timeout.</param>
        /// <returns>A task that represents the completion of the supplied task.</returns>
        /// <exception cref="System.ArgumentNullException">If task is null.</exception>
        /// <exception cref="System.ArgumentOutOfRangeException">If <paramref name="timeout"/> is negative</exception>
        public static Task WithTimeout(this Task task, TimeSpan timeout)
        {
            if (task == null)
            {
                throw new ArgumentNullException("task");
            }

            if (timeout < TimeSpan.Zero)
            {
                throw new ArgumentOutOfRangeException("timeout", "Timeout must not be negative.");
            }

            if (task.IsCompleted)
            {
                 return task;
            }

            return AddTimeout(task, timeout);
        }

        /// <summary>
        /// Adds a timeout to your task. If task does not complete before timeout then 
        /// the task fails with a TimeoutException.
        /// </summary>
        /// <remarks>Note that the task is not automatically cancelled if the timeout occurs.</remarks>
        /// <typeparam name="T">Result type of the task.</typeparam>
        /// <param name="task">The task that should have a timeout added.</param>
        /// <param name="timeout">The timeout.</param>
        /// <returns>A task that represents the completion of the supplied task.</returns>
        /// <exception cref="System.ArgumentNullException">If task is null.</exception>
        /// <exception cref="System.ArgumentOutOfRangeException">If <paramref name="timeout"/> is negative</exception>
        public static Task<T> WithTimeout<T>(this Task<T> task, TimeSpan timeout)
        {
            if (task == null)
            {
                throw new ArgumentNullException("task");
            }

            if (timeout < TimeSpan.Zero)
            {
                throw new ArgumentOutOfRangeException("timeout", "Timeout must not be negative.");
            }

            if (task.IsCompleted)
            {
                return task;
            }

            return AddTimeout(task, timeout);
        }

        private static async Task AddTimeout(Task task, TimeSpan timeout)
        {
            var cancellationSource = new CancellationTokenSource();
            var delay = Task.Delay(timeout, cancellationSource.Token);
            await Task.WhenAny(task, delay).ConfigureAwait(false);
            if (task.IsCompleted)
            {
                cancellationSource.Cancel();
            }
            else
            {
                throw new TimeoutException();
            }
        }

        private static async Task<T> AddTimeout<T>(Task<T> task, TimeSpan timeout)
        {
            await AddTimeout((Task)task, timeout);
            return await task;
        }
    }
}
